"""
pinecone tool模块作用，主要是操作向量数据库，向量数据库的初始化、创建索引、向量数据库的查询
"""
import os
from datetime import time
from dotenv import load_dotenv
import dashscope
from http import HTTPStatus
from pinecone import Pinecone
from pinecone.models import ServerlessSpec
from langchain_community.embeddings import DashScopeEmbeddings
import logging
from typing import List, Dict, Any
import re

load_dotenv()

# 定义日志
logging.basicConfig(level=logging.INFO)  # 设定日志级别
logger = logging.getLogger(__file__)  # 拿到当前文件的日志对象


class PineconeVectorStore:
    """向量数据库操作的类"""

    def __init__(self):
        """初始化相关配置信息"""
        # 1、各种 key相关配置
        self.dashscope_api_key = os.getenv("ALIBABA_BAILIAN_API_KEY")
        self.pinecone_api_key = os.getenv("PINECONE_API_KEY")
        self.pinecone_environment = "us-east-1"
        # 2、定义其他配置
        self.index_name = "menu-item-v1"  # 索引名
        self.embedding_name = "text-embedding-v2"  # 使用的嵌入模型
        self.dimension = 1536  # 向量维度
        # 3、定义组件
        self.pc = None
        self.index = None
        self.embedding = None
        # 配置 DashScope API密钥
        if self.dashscope_api_key:
            dashscope.api_key = self.dashscope_api_key

    def initialize(self) -> bool:
        """初始化 Pinecone连接和索引"""
        try:
            if not self.pinecone_api_key:
                return False
            # 得到 Pinecone实例对象
            self.pc = Pinecone(api_key=self.pinecone_api_key)
            # 判断索引是否存在
            if not self.pc.has_index(self.index_name):
                # 创建索引
                self.pc.create_index(name=self.index_name,
                                     spec=ServerlessSpec(cloud="aws", region=self.pinecone_environment),
                                     metric="cosine",
                                     dimension=self.dimension)

            # 获取索引
            self.index = self.pc.Index(self.index_name)
            # 定义嵌入模型
            self.embedding = DashScopeEmbeddings(model=self.embedding_name, dashscope_api_key=self.dashscope_api_key)
            logger.info("Pinecone数据库初始化完毕")
            return True
        except Exception as e:
            logger.error(f"Pinecone数据库初始化完毕:{str(e)}")
            return False

    def get_embedding_from_text(self, text: str) -> List[float]:
        """文本向量化"""
        if len(text) >= 5000:
            text = text[:5000]
        # 出处: https://bailian.console.aliyun.com/?spm=5176.29619931.J__Z58Z6CX7MY__Ll8p1ZOR.1.1369521cIthQu0&tab=api#/api/?type=model&url=2712515
        try:
            resp = dashscope.TextEmbedding.call(
                model=self.embedding_name,
                input=text
            )
            if resp.status_code == HTTPStatus.OK:
                return resp["output"]["embeddings"][0]["embedding"]
            else:
                return []
        except Exception as e:
            return []

    def sync_mysql_to_pinecone(self, input_data: str = None, batch_size=5000, clean_exiting: bool = True) -> bool:
        """
        向 MySQL的数据同步到 Pinecone中(文本向量化)
        :param text: 要同步的文本内容
        :param batch_size: 同步的大小 (批量同步)
        :param clean_exiting: 同步之前是否清空 pinecone中已有数据
        :return: bool
        """
        try:
            if not self.index:
                if not self.initialize():
                    return False

            if clean_exiting:
                # 调用清空 pinecone向量数据库的方法
                self.delete_all_menu_data()

            if input_data is None:
                try:
                    from db_tool import get_all_menu_items
                    # 1、查询 MySQL获取所有的数据
                    input_data = get_all_menu_items()
                    if not input_data or input_data.startswith("查询失败") or input_data.startswith("当前没有"):
                        return False
                    # 2、处理 进行同步
                    formatted_item_lists = input_data.strip().split("\n")

                    if not formatted_item_lists:
                        return False

                    batch = []
                    for line, item_str in enumerate(formatted_item_lists, start=1):
                        clean_content = item_str.strip()
                        if not clean_content:
                            continue
                        vector = self.get_embedding_from_text(clean_content)
                        if not vector:
                            continue
                        if len(vector) != self.dimension:
                            continue

                        # 构建元数据
                        metadata = {
                            "content": clean_content,
                            "line_number": line,
                            "type": "menu-item"
                        }
                        line_number = f"line_number_{line}"
                        batch.append((line_number, vector, metadata))

                        if len(batch) >= batch_size:
                            self.index.upsert(vectors=batch)
                            batch = []
                            time.sleep(1)  # pinecone在写入时限流
                    # 处理剩余
                    if batch:
                        self.index.upsert(vectors=batch)
                        batch = []
                        time.sleep(1)  # pinecone在写入时限流

                    time.sleep(5)  # 数据都同步完毕
                    return True

                except Exception as e:
                    return False

        except Exception as e:
            return False

    def delete_all_menu_data(self) -> bool:
        """删除向量数据库的数据"""
        try:
            if not self.index:
                if not self.initialize():
                    return False
            stats = self.index.describe_index_stats()
            count = stats.total_vector_count
            if count == 0:
                return True
            self.index.delete(delete_all=True)
            return True

        except Exception as e:
            return False

    def search_similar_items(self, query: str, top_k: int = 3) -> List[Dict[str, Any]]:
        """
        根据查询文本搜索相关菜品
        Args:
            query: 查询文本
            top_k: 返回结果数量

        Returns:
            List[Dict]: 相关菜品信息列表
        """

        try:
            if not self.index:
                if not self.initialize():
                    return []
            query_vector = self.get_embedding_from_text(query)
            if not query_vector:
                return []
            result = self.index.query(
                vector=query_vector,
                top_k=top_k,
                include_metadata=True
            )
            similar_items = []
            for match in result.matches:
                format_item = {
                    "id": match.id,
                    "score": match.score,
                    "content": match.metadata.get("content"),
                    "line_number": match.metadata.get("line_number")
                }
                similar_items.append(format_item)
            return similar_items

        except Exception as e:
            return []

    def get_index_stats(self) -> Dict[str, Any]:
        """获取索引统计信息"""
        try:
            if not self.index:
                if not self.initialize():
                    return {}
            stats = self.index.describe_index_stats()
            return {
                "total_vector_count": stats.total_vector_count,
                "dimension": stats.dimension,
                "index_fullness": stats.index_fullness,
                "namespaces": dict(stats.namespaces) if stats.namespaces else {}
            }
        except Exception as e:
            return {}


# 定义全局 pinecone的全局实例
pinecone = PineconeVectorStore()


# 定义查询菜品接口
def search_menu_items(query: str, top_k: int = 2) -> List[Dict]:
    """查询菜品接口"""
    result = pinecone.search_similar_items(query, top_k)
    return [item["content"] for item in result]


# 查询菜品详细信息
def search_menu_items_with_ids(query: str, top_k: int = 3) -> Dict[str, Any]:
    """根据查询搜索相关菜品，返回详细信息包括真实菜品 ID"""
    result = pinecone.search_similar_items(query, top_k)
    real_ids = []
    for item in result:
        content = item["content"]
        match = re.search(r"菜品ID:(\d+)", content)
        if match:
            id = match.group(1)
        else:
            id = item["id"]
        real_ids.append(id)
    return {
        "contents": [item["content"] for item in result],
        "ids": real_ids,
        "scores": [item["score"] for item in result]
    }


# 测试连接
def pinecone_connect() -> bool:
    """测试 pinecone数据库连接"""
    return pinecone.initialize()


# 测试同步
def sync_from_database():
    pinecone.sync_mysql_to_pinecone()


if __name__ == "__main__":
    """
    初始化数据库
    pinecone_connect()
    """
    """
    测试同步
    sync_from_database()
    """
    """
    点菜
    dict = search_menu_items_with_ids(query="我想点清新点的菜")
    for item in dict.items():
        print(item)
    """
    dict = search_menu_items_with_ids(query="你里面一共几道菜？")
    for item in dict.items():
        print(item)
